import { EventEmitter } from 'events';
import WebSocket from 'ws';
import { 
  MCPRequest, 
  MCPResponse, 
  MCPNotification, 
  MCPServerCapabilities, 
  MCPClientConfig,
  MCPConnection 
} from '../types/mcp';

export class MCPClient extends EventEmitter {
  protected ws: WebSocket | null = null;
  protected config: MCPClientConfig;
  protected connection: MCPConnection;
  private pendingRequests: Map<string | number, { resolve: Function; reject: Function }> = new Map();
  private reconnectAttempts = 0;
  private heartbeatInterval?: NodeJS.Timeout;

  constructor(config: MCPClientConfig) {
    super();
    this.config = {
      timeout: 30000,
      retryCount: 3,
      ...config
    };
    
    this.connection = {
      id: this.generateConnectionId(),
      serverUrl: config.serverUrl,
      isConnected: false,
      lastActivity: new Date()
    };
  }

  private generateConnectionId(): string {
    return `mcp_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
  }

  async connect(): Promise<void> {
    return new Promise((resolve, reject) => {
      try {
        this.ws = new WebSocket(this.config.serverUrl);
        
        if (this.ws) {
          this.ws.on('open', () => {
            console.log(`MCP Client connected to ${this.config.serverUrl}`);
            this.connection.isConnected = true;
            this.connection.lastActivity = new Date();
            this.reconnectAttempts = 0;
            
            this.initializeConnection();
            this.startHeartbeat();
            
            this.emit('connected', this.connection);
            resolve();
          });

          this.ws.on('message', (data: WebSocket.Data) => {
            this.handleMessage(data.toString());
          });

          this.ws.on('error', (error) => {
            console.error('MCP Client error:', error);
            this.emit('error', error);
            reject(error);
          });

          this.ws.on('close', (code: number, reason: string) => {
            console.log(`MCP Client disconnected: ${code} - ${reason}`);
            this.connection.isConnected = false;
            this.stopHeartbeat();
            this.emit('disconnected', { code, reason });
            
            if (this.reconnectAttempts < this.config.retryCount!) {
              this.reconnect();
            }
          });
        }

      } catch (error) {
        reject(error);
      }
    });
  }

  private async initializeConnection(): Promise<void> {
    try {
      const capabilities = await this.sendRequest<MCPServerCapabilities>('initialize', {
        protocolVersion: '2024-11-05',
        capabilities: {
          tools: {},
          resources: {},
          prompts: {}
        }
      });
      
      this.connection.capabilities = capabilities;
      this.emit('initialized', capabilities);
    } catch (error) {
      console.error('Failed to initialize MCP connection:', error);
      throw error;
    }
  }

  private startHeartbeat(): void {
    this.heartbeatInterval = setInterval(() => {
      if (this.connection.isConnected) {
        this.sendNotification('heartbeat', { timestamp: Date.now() });
      }
    }, 30000); // 30秒间隔
  }

  private stopHeartbeat(): void {
    if (this.heartbeatInterval) {
      clearInterval(this.heartbeatInterval);
      this.heartbeatInterval = undefined;
    }
  }

  private reconnect(): void {
    this.reconnectAttempts++;
    console.log(`Attempting to reconnect (${this.reconnectAttempts}/${this.config.retryCount})...`);
    
    setTimeout(() => {
      this.connect().catch(error => {
        console.error('Reconnection failed:', error);
      });
    }, 5000 * this.reconnectAttempts); // 指数退避
  }

  private handleMessage(data: string): void {
    try {
      const message = JSON.parse(data);
      this.connection.lastActivity = new Date();

      if (this.isResponse(message)) {
        this.handleResponse(message);
      } else if (this.isNotification(message)) {
        this.handleNotification(message);
      } else {
        console.warn('Unknown message type:', message);
      }
    } catch (error) {
      console.error('Failed to parse MCP message:', error);
    }
  }

  private isResponse(message: any): message is MCPResponse {
    return message.jsonrpc === '2.0' && (message.result !== undefined || message.error !== undefined);
  }

  private isNotification(message: any): message is MCPNotification {
    return message.jsonrpc === '2.0' && message.method && !message.id;
  }

  private handleResponse(response: MCPResponse): void {
    const pending = this.pendingRequests.get(response.id);
    if (pending) {
      this.pendingRequests.delete(response.id);
      
      if (response.error) {
        pending.reject(new Error(`MCP Error ${response.error.code}: ${response.error.message}`));
      } else {
        pending.resolve(response.result);
      }
    }
  }

  private handleNotification(notification: MCPNotification): void {
    this.emit('notification', notification);
    
    switch (notification.method) {
      case 'tools_changed':
        this.refreshCapabilities();
        break;
      case 'resources_changed':
        this.refreshCapabilities();
        break;
      default:
        console.log('Received notification:', notification.method);
    }
  }

  private async refreshCapabilities(): Promise<void> {
    try {
      const capabilities = await this.sendRequest<MCPServerCapabilities>('capabilities/list');
      this.connection.capabilities = capabilities;
      this.emit('capabilities_updated', capabilities);
    } catch (error) {
      console.error('Failed to refresh capabilities:', error);
    }
  }

  async sendRequest<T>(method: string, params?: Record<string, any>): Promise<T> {
    if (!this.connection.isConnected) {
      throw new Error('MCP Client is not connected');
    }

    const id = this.generateRequestId();
    const request: MCPRequest = {
      jsonrpc: '2.0',
      id,
      method,
      params
    };

    return new Promise((resolve, reject) => {
      const timeout = setTimeout(() => {
        this.pendingRequests.delete(id);
        reject(new Error(`Request timeout: ${method}`));
      }, this.config.timeout);

      this.pendingRequests.set(id, {
        resolve: (result: T) => {
          clearTimeout(timeout);
          resolve(result);
        },
        reject: (error: Error) => {
          clearTimeout(timeout);
          reject(error);
        }
      });

      if (this.ws) {
        this.ws.send(JSON.stringify(request));
      }
    });
  }

  sendNotification(method: string, params?: Record<string, any>): void {
    if (!this.connection.isConnected) {
      throw new Error('MCP Client is not connected');
    }

    const notification: MCPNotification = {
      jsonrpc: '2.0',
      method,
      params
    };

    if (this.ws) {
      this.ws.send(JSON.stringify(notification));
    }
  }

  private generateRequestId(): string {
    return `req_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
  }

  async callTool(toolName: string, arguments_: Record<string, any>): Promise<any> {
    return this.sendRequest('tools/call', {
      name: toolName,
      arguments: arguments_
    });
  }

  async getResource(resourceUri: string): Promise<any> {
    return this.sendRequest('resources/read', {
      uri: resourceUri
    });
  }

  async listTools(): Promise<any[]> {
    return this.sendRequest('tools/list');
  }

  async listResources(): Promise<any[]> {
    return this.sendRequest('resources/list');
  }

  getConnection(): MCPConnection {
    return { ...this.connection };
  }

  isConnected(): boolean {
    return this.connection.isConnected;
  }

  async disconnect(): Promise<void> {
    this.stopHeartbeat();
    
    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }
    
    this.connection.isConnected = false;
    this.pendingRequests.clear();
    this.emit('disconnected');
  }
}